import os, torch

MAP_LOC = os.environ.get('TORCH_LOAD_MAP_LOCATION', 'cpu')
_orig_torch_load = torch.load
def patched_torch_load(*args, **kwargs):
    if 'map_location' not in kwargs:
        kwargs['map_location'] = MAP_LOC
    return _orig_torch_load(*args, **kwargs)
torch.load = patched_torch_load